{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "initial_id",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:07.967505Z",
     "start_time": "2025-05-07T08:17:05.192666Z"
    },
    "collapsed": true
   },
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "D:\\conda\\a311\\Lib\\site-packages\\tqdm\\auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    }
   ],
   "source": [
    "import itertools\n",
    "import re\n",
    "import json\n",
    "import jsonlines\n",
    "import psutil\n",
    "import ujson\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "from transformers import AutoTokenizer\n",
    "from datasets import load_dataset\n",
    "import os\n",
    "from tqdm import tqdm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "c033ff9a434bead2",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:07.973788Z",
     "start_time": "2025-05-07T08:17:07.971062Z"
    }
   },
   "outputs": [],
   "source": [
    "bos_token = \"<s>\"\n",
    "eos_token = \"</s>\""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "a522b628",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:08.077878Z",
     "start_time": "2025-05-07T08:17:08.066492Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "加载的tokenierzer词表大小: 6400\n"
     ]
    }
   ],
   "source": [
    "tokenizer = AutoTokenizer.from_pretrained(\"model\", use_fast=False)\n",
    "print(f'加载的tokenierzer词表大小: {len(tokenizer)}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "9648e262",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:09.277269Z",
     "start_time": "2025-05-07T08:17:08.152170Z"
    }
   },
   "outputs": [],
   "source": [
    "import torch\n",
    "from LMConfig import LMConfig\n",
    "from model import Transformer"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "4c6bc96a",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:48.677767Z",
     "start_time": "2025-05-07T08:17:47.696739Z"
    }
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Transformer(\n",
      "  (tok_embeddings): Embedding(6400, 512)\n",
      "  (dropout): Dropout(p=0.0, inplace=False)\n",
      "  (layers): ModuleList(\n",
      "    (0-7): 8 x TransformerBlock(\n",
      "      (attention): Attention(\n",
      "        (wq): Linear(in_features=512, out_features=512, bias=False)\n",
      "        (wk): Linear(in_features=512, out_features=256, bias=False)\n",
      "        (wv): Linear(in_features=512, out_features=256, bias=False)\n",
      "        (wo): Linear(in_features=512, out_features=512, bias=False)\n",
      "        (attn_dropout): Dropout(p=0.0, inplace=False)\n",
      "        (resid_dropout): Dropout(p=0.0, inplace=False)\n",
      "      )\n",
      "      (attention_norm): RMSNorm()\n",
      "      (fnn_norm): RMSNorm()\n",
      "      (feed_forward): MOEFeedForward(\n",
      "        (experts): ModuleList(\n",
      "          (0-3): 4 x FeedForward(\n",
      "            (w1): Linear(in_features=512, out_features=1408, bias=False)\n",
      "            (w2): Linear(in_features=1408, out_features=512, bias=False)\n",
      "            (w3): Linear(in_features=512, out_features=1408, bias=False)\n",
      "            (dropout): Dropout(p=0.0, inplace=False)\n",
      "          )\n",
      "        )\n",
      "        (gate): MOEGate()\n",
      "        (shared_experts): FeedForward(\n",
      "          (w1): Linear(in_features=512, out_features=1408, bias=False)\n",
      "          (w2): Linear(in_features=1408, out_features=512, bias=False)\n",
      "          (w3): Linear(in_features=512, out_features=1408, bias=False)\n",
      "          (dropout): Dropout(p=0.0, inplace=False)\n",
      "        )\n",
      "      )\n",
      "    )\n",
      "  )\n",
      "  (norm): RMSNorm()\n",
      "  (output): Linear(in_features=512, out_features=6400, bias=False)\n",
      ")\n"
     ]
    }
   ],
   "source": [
    "config = LMConfig()\n",
    "model = Transformer(config)\n",
    "device = torch.device(\"cuda\" if torch.cuda.is_available() else \"cpu\")\n",
    "model.to(device=device)\n",
    "print(model)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "0a490960",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:49.858096Z",
     "start_time": "2025-05-07T08:17:49.601505Z"
    }
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "Transformer(\n",
       "  (tok_embeddings): Embedding(6400, 512)\n",
       "  (dropout): Dropout(p=0.0, inplace=False)\n",
       "  (layers): ModuleList(\n",
       "    (0-7): 8 x TransformerBlock(\n",
       "      (attention): Attention(\n",
       "        (wq): Linear(in_features=512, out_features=512, bias=False)\n",
       "        (wk): Linear(in_features=512, out_features=256, bias=False)\n",
       "        (wv): Linear(in_features=512, out_features=256, bias=False)\n",
       "        (wo): Linear(in_features=512, out_features=512, bias=False)\n",
       "        (attn_dropout): Dropout(p=0.0, inplace=False)\n",
       "        (resid_dropout): Dropout(p=0.0, inplace=False)\n",
       "      )\n",
       "      (attention_norm): RMSNorm()\n",
       "      (fnn_norm): RMSNorm()\n",
       "      (feed_forward): MOEFeedForward(\n",
       "        (experts): ModuleList(\n",
       "          (0-3): 4 x FeedForward(\n",
       "            (w1): Linear(in_features=512, out_features=1408, bias=False)\n",
       "            (w2): Linear(in_features=1408, out_features=512, bias=False)\n",
       "            (w3): Linear(in_features=512, out_features=1408, bias=False)\n",
       "            (dropout): Dropout(p=0.0, inplace=False)\n",
       "          )\n",
       "        )\n",
       "        (gate): MOEGate()\n",
       "        (shared_experts): FeedForward(\n",
       "          (w1): Linear(in_features=512, out_features=1408, bias=False)\n",
       "          (w2): Linear(in_features=1408, out_features=512, bias=False)\n",
       "          (w3): Linear(in_features=512, out_features=1408, bias=False)\n",
       "          (dropout): Dropout(p=0.0, inplace=False)\n",
       "        )\n",
       "      )\n",
       "    )\n",
       "  )\n",
       "  (norm): RMSNorm()\n",
       "  (output): Linear(in_features=512, out_features=6400, bias=False)\n",
       ")"
      ]
     },
     "execution_count": 9,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "model.load_state_dict(torch.load(\"out/pretrain_512_moe.pth\", map_location=\"cuda:0\"), strict=False)\n",
    "model.eval()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "1ec8eaea21b18b83",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:17:51.419223Z",
     "start_time": "2025-05-07T08:17:51.413315Z"
    }
   },
   "outputs": [],
   "source": [
    "# 准备输入文本\n",
    "input_text = \"长江是\"\n",
    "input_ids = tokenizer.encode(input_text, return_tensors='pt').to(device)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "a60a530edf4556c9",
   "metadata": {
    "ExecuteTime": {
     "end_time": "2025-05-07T08:20:16.154448Z",
     "start_time": "2025-05-07T08:20:13.697841Z"
    }
   },
   "outputs": [
    {
     "ename": "AssertionError",
     "evalue": "",
     "output_type": "error",
     "traceback": [
      "\u001b[31m---------------------------------------------------------------------------\u001b[39m",
      "\u001b[31mAssertionError\u001b[39m                            Traceback (most recent call last)",
      "\u001b[36mCell\u001b[39m\u001b[36m \u001b[39m\u001b[32mIn[11]\u001b[39m\u001b[32m, line 5\u001b[39m\n\u001b[32m      3\u001b[39m \u001b[38;5;28;01mwith\u001b[39;00m torch.no_grad():\n\u001b[32m      4\u001b[39m     \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(num_tokens_to_generate):\n\u001b[32m----> \u001b[39m\u001b[32m5\u001b[39m         output = \u001b[43mmodel\u001b[49m\u001b[43m(\u001b[49m\u001b[43minput_ids\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m      6\u001b[39m         next_token = output.logits.argmax(dim=-\u001b[32m1\u001b[39m)[:, -\u001b[32m1\u001b[39m]\n\u001b[32m      7\u001b[39m         generated_tokkens.append(next_token.item())\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\conda\\a311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1739\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1737\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m   1738\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1739\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\conda\\a311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1750\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1745\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m   1746\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m   1747\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m   1748\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m   1749\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1750\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1752\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m   1753\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\ai\\ai_leaning_all\\minids\\model.py:332\u001b[39m, in \u001b[36mTransformer.forward\u001b[39m\u001b[34m(self, tokens, targets, kv_cache, **keyargs)\u001b[39m\n\u001b[32m    329\u001b[39m pos_cis = \u001b[38;5;28mself\u001b[39m.pos_cis[current_idx:current_idx + seq_len]\n\u001b[32m    331\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m idx, layer \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(\u001b[38;5;28mself\u001b[39m.layers):\n\u001b[32m--> \u001b[39m\u001b[32m332\u001b[39m     h = \u001b[43mlayer\u001b[49m\u001b[43m(\u001b[49m\u001b[43mh\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpos_cis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkv_cache\u001b[49m\u001b[43m=\u001b[49m\u001b[43mkv_cache\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    334\u001b[39m h = \u001b[38;5;28mself\u001b[39m.norm(h)\n\u001b[32m    336\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m targets \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m:\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\conda\\a311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1739\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1737\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m   1738\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1739\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\conda\\a311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1750\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1745\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m   1746\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m   1747\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m   1748\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m   1749\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1750\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1752\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m   1753\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\ai\\ai_leaning_all\\minids\\model.py:281\u001b[39m, in \u001b[36mTransformerBlock.forward\u001b[39m\u001b[34m(self, x, pos_cis, kv_cache)\u001b[39m\n\u001b[32m    280\u001b[39m \u001b[38;5;28;01mdef\u001b[39;00m\u001b[38;5;250m \u001b[39m\u001b[34mforward\u001b[39m(\u001b[38;5;28mself\u001b[39m, x: torch.Tensor, pos_cis, kv_cache=\u001b[38;5;28;01mFalse\u001b[39;00m):\n\u001b[32m--> \u001b[39m\u001b[32m281\u001b[39m     h = x + \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mattention\u001b[49m\u001b[43m(\u001b[49m\u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43mattention_norm\u001b[49m\u001b[43m(\u001b[49m\u001b[43mx\u001b[49m\u001b[43m)\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpos_cis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mkv_cache\u001b[49m\u001b[43m=\u001b[49m\u001b[43mkv_cache\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m    282\u001b[39m     out = h + \u001b[38;5;28mself\u001b[39m.feed_forward(\u001b[38;5;28mself\u001b[39m.fnn_norm(h))\n\u001b[32m    283\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m out\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\conda\\a311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1739\u001b[39m, in \u001b[36mModule._wrapped_call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1737\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28mself\u001b[39m._compiled_call_impl(*args, **kwargs)  \u001b[38;5;66;03m# type: ignore[misc]\u001b[39;00m\n\u001b[32m   1738\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m:\n\u001b[32m-> \u001b[39m\u001b[32m1739\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[38;5;28;43mself\u001b[39;49m\u001b[43m.\u001b[49m\u001b[43m_call_impl\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\conda\\a311\\Lib\\site-packages\\torch\\nn\\modules\\module.py:1750\u001b[39m, in \u001b[36mModule._call_impl\u001b[39m\u001b[34m(self, *args, **kwargs)\u001b[39m\n\u001b[32m   1745\u001b[39m \u001b[38;5;66;03m# If we don't have any hooks, we want to skip the rest of the logic in\u001b[39;00m\n\u001b[32m   1746\u001b[39m \u001b[38;5;66;03m# this function, and just call forward.\u001b[39;00m\n\u001b[32m   1747\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m (\u001b[38;5;28mself\u001b[39m._backward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_hooks \u001b[38;5;129;01mor\u001b[39;00m \u001b[38;5;28mself\u001b[39m._forward_pre_hooks\n\u001b[32m   1748\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_backward_pre_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_backward_hooks\n\u001b[32m   1749\u001b[39m         \u001b[38;5;129;01mor\u001b[39;00m _global_forward_hooks \u001b[38;5;129;01mor\u001b[39;00m _global_forward_pre_hooks):\n\u001b[32m-> \u001b[39m\u001b[32m1750\u001b[39m     \u001b[38;5;28;01mreturn\u001b[39;00m \u001b[43mforward_call\u001b[49m\u001b[43m(\u001b[49m\u001b[43m*\u001b[49m\u001b[43margs\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43m*\u001b[49m\u001b[43m*\u001b[49m\u001b[43mkwargs\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m   1752\u001b[39m result = \u001b[38;5;28;01mNone\u001b[39;00m\n\u001b[32m   1753\u001b[39m called_always_called_hooks = \u001b[38;5;28mset\u001b[39m()\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\ai\\ai_leaning_all\\minids\\model.py:97\u001b[39m, in \u001b[36mAttention.forward\u001b[39m\u001b[34m(self, x, pos_cis, kv_cache)\u001b[39m\n\u001b[32m     94\u001b[39m xk = xk.view(bsz, seq_len, \u001b[38;5;28mself\u001b[39m.n_kv_heads, \u001b[38;5;28mself\u001b[39m.head_dim)\n\u001b[32m     95\u001b[39m xv = xv.view(bsz, seq_len, \u001b[38;5;28mself\u001b[39m.n_kv_heads, \u001b[38;5;28mself\u001b[39m.head_dim)\n\u001b[32m---> \u001b[39m\u001b[32m97\u001b[39m xq, xk = \u001b[43mapply_rotary_emb\u001b[49m\u001b[43m(\u001b[49m\u001b[43mxq\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mxk\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mpos_cis\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     99\u001b[39m \u001b[38;5;28;01mif\u001b[39;00m kv_cache \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mself\u001b[39m.eval():\n\u001b[32m    100\u001b[39m     \u001b[38;5;28;01mif\u001b[39;00m seq_len == \u001b[32m1\u001b[39m \u001b[38;5;129;01mand\u001b[39;00m \u001b[38;5;28mall\u001b[39m(cache \u001b[38;5;129;01mis\u001b[39;00m \u001b[38;5;129;01mnot\u001b[39;00m \u001b[38;5;28;01mNone\u001b[39;00m \u001b[38;5;28;01mfor\u001b[39;00m cache \u001b[38;5;129;01min\u001b[39;00m (\u001b[38;5;28mself\u001b[39m.k_cache, \u001b[38;5;28mself\u001b[39m.v_cache)):\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\ai\\ai_leaning_all\\minids\\model.py:48\u001b[39m, in \u001b[36mapply_rotary_emb\u001b[39m\u001b[34m(xq, xk, pos_cis)\u001b[39m\n\u001b[32m     46\u001b[39m xq_ = torch.view_as_complex(xq.float().reshape(*xq.shape[:-\u001b[32m1\u001b[39m], -\u001b[32m1\u001b[39m, \u001b[32m2\u001b[39m))\n\u001b[32m     47\u001b[39m xk_ = torch.view_as_complex(xk.float().reshape(*xk.shape[:-\u001b[32m1\u001b[39m], -\u001b[32m1\u001b[39m, \u001b[32m2\u001b[39m))\n\u001b[32m---> \u001b[39m\u001b[32m48\u001b[39m pos_cis = \u001b[43munite_shape\u001b[49m\u001b[43m(\u001b[49m\u001b[43mpos_cis\u001b[49m\u001b[43m,\u001b[49m\u001b[43m \u001b[49m\u001b[43mxq_\u001b[49m\u001b[43m)\u001b[49m\n\u001b[32m     49\u001b[39m xq_out = torch.view_as_real(xq_ * pos_cis).flatten(\u001b[32m3\u001b[39m)\n\u001b[32m     50\u001b[39m xk_out = torch.view_as_real(xk_ * pos_cis).flatten(\u001b[32m3\u001b[39m)\n",
      "\u001b[36mFile \u001b[39m\u001b[32mD:\\ai\\ai_leaning_all\\minids\\model.py:42\u001b[39m, in \u001b[36mapply_rotary_emb.<locals>.unite_shape\u001b[39m\u001b[34m(pos_cis, x)\u001b[39m\n\u001b[32m     40\u001b[39m ndim = x.ndim\n\u001b[32m     41\u001b[39m \u001b[38;5;28;01massert\u001b[39;00m \u001b[32m0\u001b[39m <= \u001b[32m1\u001b[39m < ndim\n\u001b[32m---> \u001b[39m\u001b[32m42\u001b[39m \u001b[38;5;28;01massert\u001b[39;00m pos_cis.shape == (x.shape[\u001b[32m1\u001b[39m], x.shape[-\u001b[32m1\u001b[39m])\n\u001b[32m     43\u001b[39m shape = [d \u001b[38;5;28;01mif\u001b[39;00m i == \u001b[32m1\u001b[39m \u001b[38;5;129;01mor\u001b[39;00m i == ndim - \u001b[32m1\u001b[39m \u001b[38;5;28;01melse\u001b[39;00m \u001b[32m1\u001b[39m \u001b[38;5;28;01mfor\u001b[39;00m i, d \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28menumerate\u001b[39m(x.shape)]\n\u001b[32m     44\u001b[39m \u001b[38;5;28;01mreturn\u001b[39;00m pos_cis.view(*shape)\n",
      "\u001b[31mAssertionError\u001b[39m: "
     ]
    }
   ],
   "source": [
    "num_tokens_to_generate = 512\n",
    "generated_tokkens = []\n",
    "with torch.no_grad():\n",
    "    for i in range(num_tokens_to_generate):\n",
    "        print(tokenizer.decode(input_ids[0]))\n",
    "        output = model(input_ids)\n",
    "        next_token = output.logits.argmax(dim=-1)[:, -1]\n",
    "        generated_tokkens.append(next_token.item())\n",
    "        input_ids = torch.cat([input_ids, next_token.unsqueeze(0)], dim=1)\n",
    "\n",
    "generated_text = tokenizer.decoder(generated_tokkens, skip_special_tokens=True)\n",
    "print(generated_text)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a9f52f091476c5d9",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "a311",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.11.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
